﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Windows.Forms;
using System.Drawing;
using System.Diagnostics;

namespace SynoManager.GUI
{
    public class TreeViewEx : TreeView
    {
        internal ImageList imageList = new ImageList();
        internal Dictionary<Image, int> imageIndices = new Dictionary<Image, int>();
        internal Dictionary<TreeNode, TreeViewEventHandler> nodeClickHandlers = new Dictionary<TreeNode, TreeViewEventHandler>();

        protected override void OnAfterSelect(TreeViewEventArgs e)
        {
            base.OnAfterSelect(e);

            if (nodeClickHandlers.ContainsKey(e.Node))
            {
                var handler = nodeClickHandlers[e.Node];
                if (handler != null)
                    handler(this, e);
            }
        }
    }

    public static class TreeViewExtensions
    {
        public static TreeNode Add(this TreeNodeCollection collection, string text, Image image, object tag)
        {
            return Add(collection, text, image, null, tag);
        }

        public static TreeNode Add(this TreeNodeCollection collection, string text, Image image, TreeViewEventHandler onSelect, object tag)
        {
            var node = collection.Add(text);
            node.Tag = tag;
            var treeView = node.TreeView as TreeViewEx;

            Debug.Assert(treeView != null, "Can only use extended add method on TreeViewEx");

            if (image != null)
            {
                int index;
                if (treeView.imageIndices.ContainsKey(image))
                {
                    index = treeView.imageIndices[image];
                }
                else
                {
                    treeView.imageList.Images.Add(image);
                    index = treeView.imageList.Images.Count - 1;
                    treeView.imageIndices.Add(image, index);
                }
                node.ImageIndex = node.SelectedImageIndex = index;
            }

            if (onSelect != null)
                treeView.nodeClickHandlers.Add(node, onSelect);

            return node;
        }
    }
}
